#include "data-collector.h"
#include "ns3/log.h"
#include "ns3/simulator.h"
#include <fstream>
#include <iomanip>
#include <sys/stat.h>
#include <sys/types.h>

namespace ns3 {

NS_LOG_COMPONENT_DEFINE("SAGINIoVDataCollector");

DataCollector::DataCollector() 
    : m_realTimeLoggingEnabled(false),
      m_loggingFrequency(1),
      m_lastLogTime(Seconds(0)),
      m_collectThroughput(false),
      m_collectDelay(false),
      m_collectInterference(false),
      m_collectPathLoss(false)
{
    NS_LOG_FUNCTION(this);
}

DataCollector::~DataCollector()
{
    NS_LOG_FUNCTION(this);
    
    // 关闭所有打开的文件
    for (auto& [vehicleId, file] : m_vehicleOutputFiles) {
        if (file.is_open()) {
            file.close();
        }
    }
}

void DataCollector::SetSimulationTimeStep(double timeStep)
{
    NS_LOG_FUNCTION(this << timeStep);
    m_timeStep = timeStep;
}

void DataCollector::SetOutputDirectory(const std::string& directory)
{
    NS_LOG_FUNCTION(this << directory);
    m_outputDirectory = directory;
}

void DataCollector::EnableThroughputCollection()
{
    NS_LOG_FUNCTION(this);
    m_collectThroughput = true;
}

void DataCollector::EnableDelayCollection()
{
    NS_LOG_FUNCTION(this);
    m_collectDelay = true;
}

void DataCollector::EnableInterferenceCollection()
{
    NS_LOG_FUNCTION(this);
    m_collectInterference = true;
}

void DataCollector::EnablePathLossCollection()
{
    NS_LOG_FUNCTION(this);
    m_collectPathLoss = true;
}

void DataCollector::RecordAllocationResult(const ChannelAllocationResult& result)
{
    NS_LOG_FUNCTION(this);
    
    // 存储当前时间点的这个车辆的结果
    m_timeVehicleResults[result.time][result.vehicleId] = result;
    
    // 添加更详细的日志，跟踪数据收集
    NS_LOG_INFO("Recorded result for vehicle " << result.vehicleId << " at time " << result.time 
              << ": Channel=" << result.channelId << ", SINR=" << result.sinr << " dB");
    
    // 定期打印当前收集的数据总数
    if (result.time - static_cast<int>(result.time) < 0.01) { // 仅在整秒时记录
        NS_LOG_INFO("Total data points collected: " << m_timeVehicleResults.size() << " time points, " 
                  << m_timeVehicleResults[result.time].size() << " vehicles at current time");
    }
}

void DataCollector::SaveResults(const std::string& filePath)
{
    NS_LOG_FUNCTION(this << filePath);
    
    // 添加日志，显示存储的数据量
    NS_LOG_INFO("Saving " << m_timeVehicleResults.size() << " time points of data to " << filePath);
    int totalEntries = 0;
    for (const auto& [time, vehicleResults] : m_timeVehicleResults) {
        totalEntries += vehicleResults.size();
    }
    NS_LOG_INFO("Total entries to save: " << totalEntries);
    
    std::ofstream outFile(filePath);
    if (!outFile.is_open()) {
        NS_LOG_ERROR("Failed to open file: " << filePath);
        return;
    }
    
    // 写入标题行
    outFile << "Time,VehicleID,X,Y,Z,Speed,ChannelID,ConnectedNodeID,NodeType,SINR";
    if (m_collectThroughput) outFile << ",Throughput";
    if (m_collectDelay) outFile << ",Delay";
    if (m_collectInterference) outFile << ",Interference";
    if (m_collectPathLoss) outFile << ",PathLoss";
    outFile << std::endl;
    
    // 写入所有时间点的所有车辆数据
    for (const auto& [time, vehicleResults] : m_timeVehicleResults) {
        for (const auto& [vehicleId, result] : vehicleResults) {
            outFile << result.time << ","
                    << result.vehicleId << ","
                    << result.x << ","
                    << result.y << ","
                    << result.z << ","
                    << result.speed << ","
                    << result.channelId << ","
                    << result.connectedNodeId << ","
                    << (result.connectedNodeType == GROUND_STATION ? "GroundStation" : "Satellite") << ","
                    << result.sinr;
            
            if (m_collectThroughput) outFile << "," << result.throughput;
            if (m_collectDelay) outFile << "," << result.delay;
            if (m_collectInterference) outFile << "," << result.interference;
            if (m_collectPathLoss) outFile << "," << result.pathLoss;
            
            outFile << std::endl;
        }
    }
    
    outFile.close();
    NS_LOG_INFO("Aggregated results saved to " << filePath);
}

void DataCollector::EnsureOutputDirectoryExists()
{
    NS_LOG_FUNCTION(this);
    
    // 创建输出目录
    std::string command = "mkdir -p " + m_outputDirectory;
    int result = system(command.c_str());
    if (result != 0) {
        NS_LOG_ERROR("Failed to create output directory: " << m_outputDirectory);
    }
}

void DataCollector::WriteCSVHeader(std::ofstream& file)
{
    // 写入类似SUMO的标题行，但添加额外网络参数
    file << "vehicle_id,x,y,z,time,speed,channel_id,connected_node_id,node_type,sinr";
    if (m_collectThroughput) file << ",throughput";
    if (m_collectDelay) file << ",delay";
    if (m_collectInterference) file << ",interference";
    if (m_collectPathLoss) file << ",path_loss";
    file << std::endl;
}

void DataCollector::SaveVehicleResults()
{
    NS_LOG_FUNCTION(this);
    
    // 确保输出目录存在
    EnsureOutputDirectoryExists();
    
    // 按车辆分组保存数据
    std::map<std::string, std::vector<ChannelAllocationResult>> vehicleData;
    
    // 收集所有车辆的所有时间点数据
    for (const auto& [time, vehicleResults] : m_timeVehicleResults) {
        for (const auto& [vehicleId, result] : vehicleResults) {
            vehicleData[vehicleId].push_back(result);
        }
    }
    
    // 对每个车辆写入单独的文件
    for (const auto& [vehicleId, results] : vehicleData) {
        // 按时间排序结果
        std::vector<ChannelAllocationResult> sortedResults = results;
        std::sort(sortedResults.begin(), sortedResults.end(), 
            [](const ChannelAllocationResult& a, const ChannelAllocationResult& b) {
                return a.time < b.time;
            });
        
        // 创建车辆输出文件
        std::string filename = m_outputDirectory + "/vehicle_" + vehicleId + ".csv";
        std::ofstream outFile(filename);
        
        if (!outFile.is_open()) {
            NS_LOG_ERROR("Failed to open file: " << filename);
            continue;
        }
        
        // 写入标题行
        WriteCSVHeader(outFile);
        
        // 写入数据行
        for (const auto& result : sortedResults) {
            outFile << result.vehicleId << ","
                   << result.x << ","
                   << result.y << ","
                   << result.z << ","
                   << result.time << ","
                   << result.speed << ","
                   << result.channelId << ","
                   << result.connectedNodeId << ","
                   << (result.connectedNodeType == GROUND_STATION ? "GroundStation" : "Satellite") << ","
                   << result.sinr;
            
            if (m_collectThroughput) outFile << "," << result.throughput;
            if (m_collectDelay) outFile << "," << result.delay;
            if (m_collectInterference) outFile << "," << result.interference;
            if (m_collectPathLoss) outFile << "," << result.pathLoss;
            
            outFile << std::endl;
        }
        
        outFile.close();
        NS_LOG_INFO("Vehicle " << vehicleId << " results saved to " << filename);
    }
    
    NS_LOG_INFO("All vehicle results saved to " << m_outputDirectory);
}

void DataCollector::CollectDataAtTime(double time)
{
    NS_LOG_FUNCTION(this << time);
    
    // 记录下个时间点的事件
    if (time < Simulator::GetMaximumSimulationTime().GetSeconds()) {
        Simulator::Schedule(Seconds(m_timeStep), 
                           &DataCollector::CollectDataAtTime, 
                           this, 
                           time + m_timeStep);
    }
    
    // 这里不需要实际进行数据收集，因为这个方法会被信道分配器和main.cc调用
    // 信道分配器将调用RecordAllocationResult来记录数据
}

void DataCollector::EnableRealTimeLogging(bool enable, int frequency)
{
    m_realTimeLoggingEnabled = enable;
    m_loggingFrequency = frequency;
    
    if (m_realTimeLoggingEnabled) {
        NS_LOG_INFO("Real-time performance logging enabled with frequency: " << m_loggingFrequency << " seconds");
        ScheduleNextLog();
    }
}

void DataCollector::ScheduleNextLog()
{
    if (!m_realTimeLoggingEnabled) {
        return;
    }
    
    Time now = Simulator::Now();
    // 安排下一次日志记录时间
    if (now >= m_lastLogTime + Seconds(m_loggingFrequency)) {
        DoLogPerformance();
        m_lastLogTime = now;
    }
    
    // 安排下一次检查
    Simulator::Schedule(Seconds(m_loggingFrequency), &DataCollector::ScheduleNextLog, this);
}

void DataCollector::DoLogPerformance()
{
    Time now = Simulator::Now();
    double currentTime = now.GetSeconds();
    std::cout << "======== Performance Metrics at t=" << currentTime << "s ========" << std::endl;
    
    // 在查找数据前添加更多日志
    NS_LOG_INFO("Looking for data at time " << currentTime << "s");
    NS_LOG_INFO("Total time points available: " << m_timeVehicleResults.size());
    
    // 输出所有可用的时间点
    if (m_timeVehicleResults.empty()) {
        NS_LOG_WARN("No time points are available in the collected data!");
    } else {
        std::string availableTimes = "Available time points: ";
        int count = 0;
        for (const auto& [time, _] : m_timeVehicleResults) {
            if (count < 10) { // 限制输出前10个时间点
                availableTimes += std::to_string(time) + " ";
                count++;
            } else {
                availableTimes += "...";
                break;
            }
        }
        NS_LOG_INFO(availableTimes);
    }
    
    // 获取当前时间点的结果
    auto timeIt = m_timeVehicleResults.find(currentTime);
    if (timeIt == m_timeVehicleResults.end()) {
        std::cout << "No data available for current time point." << std::endl;
        std::cout << "=====================================================" << std::endl;
        
        // 添加更多调试信息
        NS_LOG_WARN("No data found for time " << currentTime);
        return;
    }
    
    const auto& vehicleResults = timeIt->second;
    int totalVehicles = vehicleResults.size();
    
    if (totalVehicles == 0) {
        std::cout << "No active vehicles at current time point." << std::endl;
        std::cout << "=====================================================" << std::endl;
        return;
    }
    
    std::cout << "Total active vehicles: " << totalVehicles << std::endl;
    
    // 统计连接到地面站和卫星的车辆数量
    int groundStationConnections = 0;
    int satelliteConnections = 0;
    
    // 显示吞吐量信息
    if (m_collectThroughput) {
        double totalThroughput = 0.0;
        double maxThroughput = 0.0;
        double minThroughput = std::numeric_limits<double>::max();
        
        for (const auto& [vehicleId, result] : vehicleResults) {
            totalThroughput += result.throughput;
            maxThroughput = std::max(maxThroughput, result.throughput);
            minThroughput = std::min(minThroughput, result.throughput);
            
            if (result.connectedNodeType == GROUND_STATION) {
                groundStationConnections++;
            } else {
                satelliteConnections++;
            }
        }
        
        double avgThroughput = totalThroughput / totalVehicles;
        std::cout << "Throughput - Avg: " << std::fixed << std::setprecision(2) 
                  << avgThroughput << " Mbps, Min: " << minThroughput 
                  << " Mbps, Max: " << maxThroughput << " Mbps" << std::endl;
    }
    
    // 显示延迟信息
    if (m_collectDelay) {
        double totalDelay = 0.0;
        double maxDelay = 0.0;
        double minDelay = std::numeric_limits<double>::max();
        
        for (const auto& [vehicleId, result] : vehicleResults) {
            totalDelay += result.delay;
            maxDelay = std::max(maxDelay, result.delay);
            minDelay = std::min(minDelay, result.delay);
        }
        
        double avgDelay = totalDelay / totalVehicles;
        std::cout << "Delay - Avg: " << std::fixed << std::setprecision(2) 
                  << avgDelay << " ms, Min: " << minDelay 
                  << " ms, Max: " << maxDelay << " ms" << std::endl;
    }
    
    // 显示SINR信息
    double totalSINR = 0.0;
    double maxSINR = -std::numeric_limits<double>::max();
    double minSINR = std::numeric_limits<double>::max();
    
    for (const auto& [vehicleId, result] : vehicleResults) {
        totalSINR += result.sinr;
        maxSINR = std::max(maxSINR, result.sinr);
        minSINR = std::min(minSINR, result.sinr);
    }
    
    double avgSINR = totalSINR / totalVehicles;
    std::cout << "SINR - Avg: " << std::fixed << std::setprecision(2) 
              << avgSINR << " dB, Min: " << minSINR 
              << " dB, Max: " << maxSINR << " dB" << std::endl;
    
    // 显示干扰信息
    if (m_collectInterference) {
        double totalInterference = 0.0;
        double maxInterference = 0.0;
        double minInterference = std::numeric_limits<double>::max();
        
        for (const auto& [vehicleId, result] : vehicleResults) {
            totalInterference += result.interference;
            maxInterference = std::max(maxInterference, result.interference);
            minInterference = std::min(minInterference, result.interference);
        }
        
        double avgInterference = totalInterference / totalVehicles;
        std::cout << "Interference - Avg: " << std::fixed << std::setprecision(2) 
                  << avgInterference << " dBm, Min: " << minInterference 
                  << " dBm, Max: " << maxInterference << " dBm" << std::endl;
    }
    
    // 显示路径损耗信息
    if (m_collectPathLoss) {
        double totalPathLoss = 0.0;
        double maxPathLoss = 0.0;
        double minPathLoss = std::numeric_limits<double>::max();
        
        for (const auto& [vehicleId, result] : vehicleResults) {
            totalPathLoss += result.pathLoss;
            maxPathLoss = std::max(maxPathLoss, result.pathLoss);
            minPathLoss = std::min(minPathLoss, result.pathLoss);
        }
        
        double avgPathLoss = totalPathLoss / totalVehicles;
        std::cout << "Path Loss - Avg: " << std::fixed << std::setprecision(2) 
                  << avgPathLoss << " dB, Min: " << minPathLoss 
                  << " dB, Max: " << maxPathLoss << " dB" << std::endl;
    }
    
    // 显示连接情况
    std::cout << "Connection Distribution - Ground Stations: " << groundStationConnections 
              << " (" << std::fixed << std::setprecision(1) 
              << (groundStationConnections * 100.0 / totalVehicles) << "%), "
              << "Satellites: " << satelliteConnections 
              << " (" << (satelliteConnections * 100.0 / totalVehicles) << "%)" 
              << std::endl;
    
    std::cout << "=====================================================" << std::endl;
}

void DataCollector::StartPeriodicCollection(double startTime, double endTime)
{
    NS_LOG_FUNCTION(this << startTime << endTime);
    
    // 设置模拟器最大时间
    Simulator::Stop(Seconds(endTime));
    
    // 安排第一次数据收集
    Simulator::Schedule(Seconds(startTime), 
                       &DataCollector::CollectDataAtTime, 
                       this, 
                       startTime);
    
    // 如果启用了实时日志，设置第一次日志记录
    if (m_realTimeLoggingEnabled) {
        Simulator::Schedule(Seconds(startTime), &DataCollector::ScheduleNextLog, this);
    }
    
    NS_LOG_INFO("Periodic data collection scheduled from " << startTime 
               << "s to " << endTime << "s with step " << m_timeStep << "s");
}

} // namespace ns3
